def compute_single_result(op_name, num_args):
  """An implementation of the most popular ops from the MathQA dataset."""
  # See https://gitlab.cs.washington.edu/amini91/mathqa-categorization/
  # and specfically line 142 and following in new_DataStructure.py
  # for an implementation which covers more details.
  if op_name == 'add':
    return num_args[0] + num_args[1]
  elif op_name == 'circle_arc':
    return num_args[0] / 360 * math.pi * 2 * num_args[1]
  elif op_name == 'circle_area':
    return math.pi * num_args[0]**2
  elif op_name == 'circle_sector_area':
    return num_args[1] / 360 * math.pi * (num_args[0]**2)
  elif op_name == 'circumface':
    return 2 * math.pi * num_args[0]
  elif op_name == 'choose':
    return scipy.special.comb(num_args[0], num_args[1])
  elif op_name == 'cosine':
    return math.cos(num_args[0])
  elif op_name == 'cube_edge_by_volume':
    return num_args[0]**(1 / 3)
  elif op_name == 'combined_work':
    return 1 / (
        min(num_args[0], 1 / num_args[0]) + min(num_args[1], 1 / num_args[1]))
  elif op_name == 'count_interval':
    return num_args[0] - num_args[1] + 1
  elif op_name == 'diagonal':
    return math.sqrt(num_args[0]**2 + num_args[1]**2)
  elif op_name == 'divide' or op_name == 'speed':
    if num_args[1] != 0:
      return num_args[0] / num_args[1]
    else:
      return 0
  elif op_name == 'factorial':
    return math.factorial(min(15, int(num_args[0])))
  elif op_name == 'floor':
    return math.floor(num_args[0])
  elif op_name == 'find_work':
    return 1 / (
        max(
            min(num_args[0], 1 / num_args[0]), min(
                num_args[1], 1 / num_args[1])) - min(
                    min(num_args[0], 1 / num_args[0]),
                    min(num_args[1], 1 / num_args[1])))
  elif op_name == 'from_percent':
    return num_args[0] / 100
  elif op_name == 'gain_percent':
    return 100 + num_args[0]
  elif op_name == 'gcd':
    return np.gcd(int(num_args[0]), int(num_args[1]))
  elif op_name == 'inverse':
    if num_args[0] != 0:
      return 1 / num_args[0]
    else:
      return 0
  elif op_name == 'lcm':
    return scipy.lcm(int(num_args[0]), int(num_args[1]))
  elif op_name == 'log':
    return math.log(max(1e-5, num_args[0]), 2)
  elif op_name == 'loss_percent':
    return 100 - num_args[0]
  elif op_name == 'max':
    return max(num_args[0], num_args[1])
  elif op_name == 'multiply':
    return num_args[0] * num_args[1]
  elif op_name == 'negate_percent':
    return 100 - num_args[0]
  elif op_name == 'negate':
    return -num_args[0]
  elif op_name == 'original_price_before_loss':
    return num_args[1] * 100 / (100 + 1e-5 - num_args[0])
  elif op_name == 'original_price_before_gain':
    return num_args[1] * 100 / (100 + num_args[0])
  elif op_name == 'permutation':
    n, m = min(num_args[0], num_args[1]), max(num_args[0], num_args[1])
    return math.factorial(int(m)) / math.factorial(int(m - n))
  elif op_name == 'power':
    return num_args[0]**min(num_args[1], 5)
  elif op_name == 'percent':
    return num_args[0] / 100 * num_args[1]
  elif op_name == 'price_after_gain' or op_name == 'p_after_gain':
    return (1 + num_args[0] / 100) * num_args[1]
  elif op_name == 'price_after_loss' or op_name == 'price_after_loss':
    return (1 - num_args[0] / 100) * num_args[1]
  elif op_name == 'quadrilateral_area':
    return num_args[0] * (num_args[1] + num_args[2]) / 2
  elif op_name == 'reminder':
    return num_args[0] % num_args[1]
  elif op_name == 'rectangle_area':
    return num_args[0] * num_args[1]
  elif op_name == 'rectangle_perimeter':
    return 2 * (num_args[0] + num_args[1])
  elif op_name == 'rhombus_area':
    return num_args[0] * num_args[1] / 2
  elif op_name == 'sine':
    return math.sin(num_args[0])
  elif op_name == 'sqrt':
    return math.sqrt(max(0, num_args[0]))
  elif op_name == 'subtract':
    return num_args[0] - num_args[1]
  elif op_name == 'square_edge_by_perimeter':
    return num_args[0] / 4
  elif op_name == 'square_edge_by_area':
    return math.sqrt(num_args[0])
  elif op_name == 'square_area':
    return num_args[0]**2
  elif op_name == 'surface_cube':
    return 6 * num_args[0]**2
  elif op_name == 'surface_rectangular_prism':
    return 2 * (
        num_args[0] * num_args[1] + num_args[0] * num_args[2] +
        num_args[1] * num_args[2])
  elif op_name == 'semi_circle_perimiter':
    return math.pi * num_args[0] + 2 * num_args[0]
  elif op_name == 'square_perimeter' or op_name == 'rhombus_perimeter':
    return 4 * num_args[0]
  elif op_name == 'surface_sphere':
    return 4 * math.pi * num_args[0]**2
  elif op_name == 'speed_ratio_steel_to_stream':
    return (num_args[0] + num_args[1]) / (num_args[0] - num_args[1])
  elif op_name == 'speed_in_still_water':
    return (num_args[0] + num_args[1]) / 2
  elif op_name == 'stream_speed':
    return (num_args[0] - num_args[1]) / 2
  elif op_name == 'trapezium_area':
    return num_args[0] * (num_args[1] + num_args[2]) / 2
  elif op_name == 'triangle_area':
    return num_args[0] * num_args[1] / 2
  elif op_name == 'triangle_perimeter':
    return num_args[0] + num_args[1] + num_args[2]
  elif op_name == 'triangle_area_three_edges':
    # Heron's formula
    s = (num_args[0] + num_args[1] + num_args[2]) / 2
    return math.sqrt(
        max(0,
            s * (s - num_args[0]) * (s - num_args[1]) * (s - num_args[2])))
  elif op_name == 'union_prob':
    return num_args[0] + num_args[1] - num_args[2]
  elif op_name == 'negate_prob':
    return 1 - num_args[0]
  elif op_name == 'volume_cube':
    return num_args[0]**3
  elif op_name == 'volume_cone':
    return math.pi * num_args[0]**2 * num_args[1] / 3
  elif op_name == 'volume_cylinder':
    return math.pi * num_args[0]**2 * num_args[1]
  elif op_name == 'volume_rectangular_prism':
    return num_args[0] * num_args[1] * num_args[2]
  elif op_name == 'volume_sphere':
    return 4 / 3 * math.pi * num_args[0]**3
An implementation of the most popular ops from the MathQA dataset.
trax/data/tf_inputs.py