docstr: Checks support for additional dimensions in inputs.
docstr_tok: ["Checks","support","for","additional","dimensions","in","inputs","."]
code: def test_higher_rank_inputs_for_importance_weights(self):
      """Checks support for additional dimensions in inputs."""
      for fw in framework_iterator(
          frameworks=("torch", "tf"), session=True
      ):
          vtrace = vtrace_tf if fw != "torch" else vtrace_torch
          if fw == "tf":
              inputs_ = {
                  "log_rhos": tf1.placeholder(
                      dtype=tf.float32, shape=[None, None, 1]
                  ),
                  "discounts": tf1.placeholder(
                      dtype=tf.float32, shape=[None, None, 1]
                  ),
                  "rewards": tf1.placeholder(
                      dtype=tf.float32, shape=[None, None, 42]
                  ),
                  "values": tf1.placeholder(
                      dtype=tf.float32, shape=[None, None, 42]
                  ),
                  "bootstrap_value": tf1.placeholder(
                      dtype=tf.float32, shape=[None, 42]
                  ),
              }
          else:
              inputs_ = {
                  "log_rhos": Box(-1.0, 1.0, (8, 10, 1)).sample(),
                  "discounts": Box(-1.0, 1.0, (8, 10, 1)).sample(),
                  "rewards": Box(-1.0, 1.0, (8, 10, 42)).sample(),
                  "values": Box(-1.0, 1.0, (8, 10, 42)).sample(),
                  "bootstrap_value": Box(-1.0, 1.0, (10, 42)).sample(),
              }
          output = vtrace.from_importance_weights(**inputs_)
          check(int(output.vs.shape[-1]), 42)
  
file_path: ray/rllib/algorithms/impala/tests/test_vtrace.py
imports: from gymnasium.spaces import Box
  import numpy
  import unittest
  from ray.rllib.algorithms.impala import vtrace_tf
  from ray.rllib.algorithms.impala import vtrace_torch
  from ray.rllib.utils.framework import try_import_tf, try_import_torch
  from ray.rllib.utils.numpy import softmax
  from ray.rllib.utils.test_utils import check, framework_iterator
label_desc: Use image node_id 3 for calling the VtraceTest obj's underlying member method code with example usage: obj.test_higher_rank_inputs_for_importance_weights() without return types