func.func @main(%arg0: tensor<10x19xf32>, %arg1: tensor<19x10xf32> {mhlo.is_same_data_across_replicas = true}) -> tensor<10x19xf32> {
    %0 = "tf.Shape"(%arg0) : (tensor<10x19xf32>) -> tensor<2xi64>
    %1 = "tf.Reshape"(%arg1, %0) : (tensor<19x10xf32>, tensor<2xi64>) -> tensor<10x19xf32>
    func.return %1 : tensor<10x19xf32>
  }
}

// Tests that foldable ops are constant-folded to enable legalization of ops

  %1 = mhlo.convert %0 : (tensor<f64>) -> tensor<i64>
  %2 = mhlo.reshape %1 : (tensor<i64>) -> tensor<1xi64>
  // expected-error @+1 {{Node `mhlo.dynamic_iota` must have compile-time constant}}
  %3 = "mhlo.dynamic_iota"(%2) <{iota_dimension = 0 : i64}> : (tensor<1xi64>) -> tensor<?xi32>
  %4 = mhlo.multiply %3, %3 : tensor<?xi32>
  return %4 : tensor<?xi32>

    ^bb0(%arg1: tensor<f32>, %arg2: tensor<f32>):
      %2 = stablehlo.maximum %arg1, %arg2 : tensor<f32>
      stablehlo.return %2 : tensor<f32>
  }) {padding = dense<[[0, 0], [1, 1], [1, 1], [0, 0]]> : tensor<4x2xi64>, window_dimensions = array<i64: 1, 3, 3, 1>} : (tensor<2x3x1x3xf32>, tensor<f32>) -> tensor<2x3x1x3xf32>
  return %1 : tensor<2x3x1x3xf32>
}

// -----

    func.func @main() -> () {
      %cst0 = "tf.Const"(){ value = dense<0> : tensor<3x5xi1>} : () -> tensor<3x5xi1>
      %0 = "tf.Where"(%cst0) : (tensor<3x5xi1>) -> tensor<?x2xi64>
      func.return
    }
  })";

  CellReader<int64_t> dynamism_counter(kDynamismOpCounterStreamzName);
  CellReader<int64_t> dynamism_function_counter(
      kDynamismFunctionCounterStreamzName);

  ^bb0(%arg1: tensor<f32>, %arg2: tensor<f32>):
    %7 = stablehlo.maximum %arg1, %arg2 : tensor<f32>
    stablehlo.return %7 : tensor<f32>
  }) {padding = dense<[[0, 0], [0, 1], [0, 1], [0, 0]]> : tensor<4x2xi64>, window_dimensions = array<i64: 1, 3, 3, 1>, window_strides = array<i64: 1, 2, 2, 1>} : (tensor<?x112x112x64xf32>, tensor<f32>) -> tensor<?x56x56x64xf32>
  return %6 : tensor<?x56x56x64xf32>

//   [counter_lo, counter_hi, key]
//
//   RngReadAndSkip increments the 128 bit counter value by 256 * delta and
//   returns the original state value.
//
// For Threefry, the resource variable holds a tensor<2xi64> with the state:
//   [counter, key]
//
//   RngReadAndSkip increments the 64 bit counter value by 256 * delta and
//   returns a tensor<3xi64> value [counter, key, 0].
class DecomposeRngReadAndSkipOp : public RewritePattern {