docstr: Check that stateless_random_shuffle works with multi-dim shapes.
docstr_tok: ["Check","that","stateless_random_shuffle","works","with","multi-dim","shapes","."]
code: def testMultiDimensionalShape(self):
      """Check that stateless_random_shuffle works with multi-dim shapes."""
      for dtype in (tf.int32, tf.int64, tf.float32, tf.float64):
          input_permutation = tf.constant(
              [[[1], [2], [3]], [[4], [5], [6]]], dtype=dtype
          )
          random_shuffle = tff_rnd.stateless_random_shuffle(
              input_permutation, seed=(1, 42)
          )
          random_permutation_first_call = self.evaluate(random_shuffle)
          random_permutation_next_call = self.evaluate(random_shuffle)
          input_permutation = self.evaluate(input_permutation)
          # Check that the dtype is correct
          np.testing.assert_equal(
              random_permutation_first_call.dtype, dtype.as_numpy_dtype
          )
          # Check that the shuffles are the same
          np.testing.assert_array_equal(
              random_permutation_first_call,
              random_permutation_next_call,
          )
          # Check that the output shape is correct
          np.testing.assert_equal(
              random_permutation_first_call.shape,
              input_permutation.shape,
          )
  
file_path: tf-quant-finance/tf_quant_finance/math/random_ops/stateless_test.py
imports: import numpy
  import tensorflow.compat.v2
  import tf_quant_finance
  from tensorflow.python.framework import test_util
label_desc: Use image node_id 5 for calling the StatelessRandomOpsTest obj's underlying member method code with example usage: obj.testMultiDimensionalShape() without return types