christopher/tensorflow-coder

1# Based on the following code demo: https://github.com/google-research/tensorflow-coder/blob/master/tf_coder/tf_coder_main.py2import streamlit as st3from tf_coder.value_search import colab_interface, value_search_settings4from streamlit_ace import st_ace5st.set_page_config(page_title="TensorFlow Coder", page_icon='👩‍💻', layout="wide")6 7st.title("👩‍💻 TensorFlow Coder")8st.write('#')9st.write("[TensorFlow Coder](https://github.com/google-research/tensorflow-coder) is a program synthesis tool developed at Google Research by Kensen Shi, David Bieber and Rishabh Singh. It takes an example input-output tensor example and attempts to find the combination of TensorFlow ops that capture that transformation. Please cite the authors' [paper](https://github.com/google-research/tensorflow-coder/blob/master/README.md#citation) if you use their tool in your work. Also checkout the TensorFlow [Blog post](https://blog.tensorflow.org/2020/08/introducing-tensorflow-coder-tool.html) for more information and examples.")10 11col1, col2, col3 = st.columns([5, 5, 3])12with col1:13    st.write('#### Inputs')14    inputs = st_ace(placeholder="The input tensor(s) specified as a dictionary", value="{'rows': [10, 20, 30],\n'cols': [1,2,3,4]}", language="python", theme="solarized_dark", auto_update=True)15with col2:16    st.write('#### Output')17    output = st_ace(placeholder="The output tensor", value="[[11, 12, 13, 14],\n[21, 22, 23, 24],\n[31, 32, 33, 34]]", language="python", theme="solarized_dark", auto_update=True)18with col3:19    st.write('#### Constants')20    constants = st_ace(placeholder="Optional list of scalar constants", value="[]", language="python", theme="solarized_dark", auto_update=True)21 22st.write("#### Description")23description = st.text_input(label="", placeholder="An optional natural language description of the operation", value="add two vectors with broadcasting to get a matrix")24with st.expander("⚙️ Search Options", expanded=False):25    settings_kwargs = dict()26    settings_kwargs["require_all_inputs_used"] = st.checkbox("Require All Inputs", value=True)27    settings_kwargs["only_minimal_solutions"] = st.checkbox("Only Minimal Solutions", value=False)28    settings_kwargs["max_solutions"] = st.slider("Maximum number of solutions", value=1, min_value=1, step=1, max_value=256)29    settings_kwargs["timeout"] = st.slider("Timeout in seconds", value=300, min_value=1, step=10, max_value=300)30    31if st.button("🔎 Search for Tensor Ops!"):32    i = eval(inputs)33    o = eval(output)34    c = eval(constants)35    settings = value_search_settings.from_dict({36      'timeout': settings_kwargs["timeout"],37      'only_minimal_solutions': settings_kwargs["only_minimal_solutions"],38      'max_solutions': settings_kwargs["max_solutions"],39      'require_all_inputs_used': settings_kwargs["require_all_inputs_used"],40      'require_one_input_used': not settings_kwargs["require_all_inputs_used"],41  })42    with st.spinner("Searching for solution..."):43        results = colab_interface.run_value_search_from_colab(i, o, c, description, settings)44        num_solutions = len(results.solutions)45        solution_solutions = " solutions" if num_solutions > 1 else " solution"46        st.write(f"Found {num_solutions}{solution_solutions} in {results.total_time:.2f} seconds")47        for solution in results.solutions:48            st.code(solution.expression, language='python')