code: def test_shape_assertion_failure(self):
      x = np.ones((3, 5), dtype=np.int32)
      res = np.int32(x.shape[0])
  
      def f(x):  # x: f32[b, 5] and b = 3, with a constraint b == 4.
          # return x.shape[0]
          module, version = serialize(
              """
  module @jit_fun.1 attributes {jax.uses_shape_polymorphism = true} {
    func.func public @main(%arg1: tensor<?x5xi32>) -> tensor<i32> {
      %b = "stablehlo.get_dimension_size"(%arg1) {dimension = 0 : i64} : (tensor<?x5xi32>) -> tensor<i32>
      %4 = stablehlo.constant dense<4> : tensor<i32>
      %5 = stablehlo.constant dense<5> : tensor<i32>
      %11 = stablehlo.constant dense<11> : tensor<i32>
      %ok = stablehlo.compare  EQ, %b, %4,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1>
      stablehlo.custom_call @shape_assertion(%ok, %b, %4, %5, %4, %5, %4, %5, %4, %5, %4, %5, %11) {
        error_message = "Expecting {0} == {1}. Extra {2,=5}, {3}, {{0}, {4}, {5}, {6}, {7}, {11}.",
        has_side_effect = true
      } : (tensor<i1>, tensor<i32>, tensor<i32>, tensor<i32>, tensor<i32>, tensor<i32>, tensor<i32>, tensor<i32>, tensor<i32>, tensor<i32>, tensor<i32>, tensor<i32>, tensor<i32>) -> ()
      return %b : tensor<i32>
    }
  }
  """
          )
          return xla.call_module(
              [
                  x,
              ],
              version=version,
              module=module,
              Tout=[res.dtype],
              Sout=[res.shape],
              platforms=[self.testing_platform()],
          )
  
      # This test runs as part of two targets, with and without
      # disabling shape_assertions.
      disabled_shape_assertions_check = (
          "--tf_xla_call_module_disabled_checks=shape_assertions"
          in os.getenv("TF_XLA_FLAGS", "")
      )
      if disabled_shape_assertions_check:
          # No error even though the constraint is false.
          self._assertOpOutputMatchesExpected(f, (x,), (res,))
      else:
          with self.assertRaisesRegex(
              errors.InvalidArgumentError,
              re.escape(
                  "Expecting 3 == 4. Extra   5  , 4, {0}, 5, 4, 5, 4, 11."
              ),
          ):
              self._assertOpOutputMatchesExpected(f, (x,), (res,))
  
imports: import os
  import re
  from typing import Optional, Sequence
  import unittest
  from absl.testing import parameterized
  import numpy
  from tensorflow.compiler.mlir.stablehlo import stablehlo
  from tensorflow.compiler.tests import xla_test
  from tensorflow.compiler.tf2xla.ops import gen_xla_ops
  from tensorflow.compiler.tf2xla.python import xla
  from tensorflow.core.protobuf import config_pb2
  from tensorflow.python.eager import def_function
  from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import errors
  from tensorflow.python.framework import function
  from tensorflow.python.framework import ops
  from tensorflow.python.platform import googletest
  from tensorflow.python.platform import test
label_desc: Use image node_id 26 for calling the XlaCallModuleOpTest obj's underlying member method code with example usage: obj.test_shape_assertion_failure() and returns: xla