ST (r)

	CALL p256SqrInternal(SB)	// rsqr = rˆ2
	ST (rsqr)

	LDacc (h)
	CALL p256SqrInternal(SB)	// hsqr = hˆ2
	ST (hsqr)

	LDt (h)
	CALL p256MulInternal(SB)	// hcub = hˆ3
	ST (hcub)

	LDt (y1in)
	CALL p256MulInternal(SB)	// y1 * hˆ3
	ST (s2)

	LDacc (x1in)
	LDt (hsqr)
	CALL p256MulInternal(SB)	// u1 * hˆ2
	ST (h)

	STx(r)

	CALL	p256SqrInternal<>(SB)    // rsqr = rˆ2
	STy	(rsqr)

	LDx(h)
	CALL	p256SqrInternal<>(SB)    // hsqr = hˆ2
	STy(hsqr)

	CALL	p256MulInternal<>(SB)    // hcub = hˆ3
	STy(hcub)

	LDx(y1in)
	CALL	p256MulInternal<>(SB)    // y1 * hˆ3
	STy(s2)

	LDP	hsqr(0*8), (x0, x1)
	LDP	hsqr(2*8), (x2, x3)
	LDP	0*16(a_ptr), (y0, y1)
	LDP	1*16(a_ptr), (y2, y3)

func.func @main(%arg0: tensor<2xf32>) -> tensor<2xf32> {
  %0 = "tfl.custom"(%arg0) {custom_code = "stablehlo.rsqrt", custom_option = #tfl<const_bytes : "0x00000100002401">} : (tensor<2xf32>) -> tensor<2xf32>
  func.return %0 : tensor<2xf32>
}
}

// CHECK:       module
// CHECK-NEXT:  func @main(%arg0: tensor<2xf32>) -> tensor<2xf32> {
// CHECK-NEXT:  %0 = stablehlo.rsqrt %arg0 : tensor<2xf32>
// CHECK-NEXT:  return %0 : tensor<2xf32>
// CHECK-NEXT:  }

//              rsqrt(variance + epsilon)
// CHECK:  %[[RSQRT:.*]] = "tf.Rsqrt"(%[[ADD1]])
//              scale * rsqrt(variance + epsilon)
// CHECK:  %[[MUL1:.*]] = "tf.Mul"(%[[ARG1:.*]], %[[RSQRT]])
//              x * scale * rsqrt(variance + epsilon)
// CHECK:  %[[MUL2:.*]] = "tf.Mul"(%[[ARG0:.*]], %[[MUL1]])
//              mean * scale * rsqrt(variance + epsilon)

  // CHECK-DAG: %[[EPS_BCAST:.+]] = mhlo.constant dense<1.001000e-05> : tensor<256xf32>
  // CHECK-DAG: %[[VARIANCE_EPS:.+]] = mhlo.add %[[VARIANCE]], %[[EPS_BCAST]] : tensor<256xf32>
  // CHECK-DAG: %[[VARIANCE_EPS_RSQRT:.+]] = mhlo.rsqrt %[[VARIANCE_EPS]] : tensor<256xf32>
  // CHECK-DAG: %[[MULTIPLIER:.+]] = mhlo.multiply %[[VARIANCE_EPS_RSQRT]], %[[SCALE]] : tensor<256xf32>
  // CHECK-DAG: %[[MUL_MEAN:.+]] = mhlo.multiply %[[MULTIPLIER]], %[[MEAN]] : tensor<256xf32>

			// V(k'+1) = V(2k+2) = V(k+1)² - 2.
			t1 = t1.sqr(vk1)
			t1 = t1.add(t1, nm2)
			t2, vk1 = t2.div(vk1, t1, n)
		} else {
			// k' = 2k
			// V(k'+1) = V(2k+1) = V(k) V(k+1) - P.
			t1 = t1.mul(vk, vk1)
			t1 = t1.add(t1, n)
			t1 = t1.sub(t1, natP)
			t2, vk1 = t2.div(vk1, t1, n)
			// V(k') = V(2k) = V(k)² - 2
			t1 = t1.sqr(vk)
			t1 = t1.add(t1, nm2)
			t2, vk = t2.div(vk, t1, n)
		}

    $t, $m, $v, $beta, $gamma, F32Attr:$variance_epsilon, ConstBoolAttrFalse:$scale_after_normalization),
  (TF_AddOp
    (TF_MulOp $t, (TF_RsqrtOp:$rsqrt (TF_AddOp $v, (TF_ConstOp $variance_epsilon)))),
    (TF_SubOp $beta, (TF_MulOp $m, $rsqrt)))>;
def ConvertBatchNormWithGlobalNormalization_2 : Pat<
  (TF_BatchNormWithGlobalNormalizationOp

}

TEST_F(CWiseUnaryGradTest, Rsqrt) {
  auto x_fn = [this](const int i) { return RV({1, 2, 3, 4, 5, 6, 7, 8}); };
  TestCWiseGrad<float, float>(RSQRT, x_fn);
}

TEST_F(CWiseUnaryGradTest, Rsqrt_Complex) {
  auto x_fn = [this](const int i) {
    return CRV({{-1.0f, 0.5f}, {1.0f, 0.5f}, {2, -1}});
  };
  TestCWiseGrad<complex64, complex64>(RSQRT, x_fn);
}

TEST_F(CWiseUnaryGradTest, Exp) {

				// to check performance before making changes.
				zz = zz.sqr(z)
				zz, z = z, zz
				z = z.trunc(z, logM)

				zz = zz.sqr(z)
				zz, z = z, zz
				z = z.trunc(z, logM)

				zz = zz.sqr(z)
				zz, z = z, zz
				z = z.trunc(z, logM)

				zz = zz.sqr(z)
				zz, z = z, zz
				z = z.trunc(z, logM)
			}

// tfl.log / tfl.logistic / tfl.max_pool_2d / tfl.mirror_pad / tfl.maximum /
// tfl.custom / tfl.mean / tfl.minimum / tfl.pad / tfl.pow / tfl.prelu /
// tfl.relu / tfl.relu6 / tfl.rsqrt / tfl.sin / tfl.slice / tfl.softmax /
// tfl.space_to_depth / tfl.sqrt / tfl.square / tfl.squared_difference /
// tfl.strided_slice / tfl.tanh / tfl.transpose / tfl.transpose_conv