docstr: Test that we correctly handle multi-dimensional dynamic output.
docstr_tok: ["Test","that","we","correctly","handle","multi-dimensional","dynamic","output","."]
code: def test_dynamic_output_multidim(self):
      """Test that we correctly handle multi-dimensional dynamic output."""
  
      @def_function.function(jit_compile=True)
      def compiled_f(shape):
          return (
              test_ops_for_light_outside_compilation.dynamic_multidim(
                  shape
              )
          )
  
      with context.device("/gpu:0"):
          # Rank is hardcoded to 5.
          shape = [3, 4, 5, 4, 3]
          hlo = compiled_f.experimental_get_compiler_ir(shape)(
              "hlo_no_metadata"
          )
          out = compiled_f(shape)
  
          self.assertFilecheck(
              hlo,
              r"""
            CHECK: f32[<=20,<=20,<=20,<=20,<=20]{4,3,2,1,0} custom-call(), custom_call_target="GenericTfCallbackGPU"
            CHECK: DynamicMultidim
            """,
          )
          self.assertAllClose(out, array_ops.ones(shape))
  
imports: from tensorflow.compiler.mlir.python.mlir_wrapper import filecheck_wrapper
  from tensorflow.compiler.tf2xla import test_ops_for_light_outside_compilation
  from tensorflow.python.eager import context
  from tensorflow.python.eager import def_function
  from tensorflow.python.framework import ops
  from tensorflow.python.framework import test_util
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import nn_ops
  from tensorflow.python.ops import random_ops
  from tensorflow.python.platform import googletest
  from tensorflow.python.platform import test
label_desc: Use image node_id 6 for calling the LightOutsideCompilationTest obj's underlying member method code with example usage: obj.test_dynamic_output_multidim() and returns: test_ops_for_light_outside_compilation