// CHECK: "tf.IfRegion"(%[[PRED]]) <{is_stateless = false}> ({
  %case_op = "tf.IfRegion"(%arg0) <{is_stateless = false}> ({
      // CHECK: %[[TA_VAL:.*]] = "tf.ReadVariableOp"(%[[TA_BUFFER]])
      // CHECK: "tf.Slice"(%[[TA_VAL]]
      // CHECK-NOT: tf.TensorArrayReadV3
      %idx = "tf.Const"() {value = dense<3> : tensor<i32>} : () -> tensor<i32>

// Replace
//   (float)OneHot(index, depth, on_val, off_val, axis)
// With
//   OneHot(index, depth, (float)on_val, (float)off_val, axis)
def FuseOneHotAndCastToFloat : Pat<
  (TFL_CastOp:$output (TFL_OneHotOp $indices,
                                    $depth,
                                    (Arith_ConstantOp $on_val),
                                    (Arith_ConstantOp $off_val),

  def res_name(self, ns, types_ns, name):
    name_str = str(name)
    if name_str in TFR_BUILTINS:
      return {TFRTypes.TFR_BUILTIN_FUNC}, name_str
    if name_str in ns:
      ns_val = ns[name_str]
      return {type(ns_val)}, ns_val
    if name_str in __builtins__:
      return {TFRTypes.PY_BUILTIN_FUNC}, __builtins__[name_str]
    # This name is not in the namespace because the autograph transformation

          use_variable: If True, creates a variable for weight.
        """
        super(GatherModel, self).__init__()
        w_val = np.random.randn(128, 32).astype('f4')
        if use_variable:
          self.w = variables.Variable(w_val)
        else:
          self.w = w_val

      @def_function.function(
          input_signature=[
              tensor_spec.TensorSpec(