code: def test_where_partial_shape():
      x = mx.sym.Variable("x")
      y = mx.sym.Variable("y")
      cond = mx.sym.Variable("cond")
      where_op = mx.sym.where(cond, x, y)
      # condition must be fully known to infer shape
      _, result, _ = where_op.infer_shape_partial(
          cond=(0, 2), x=(0, 2), y=(0, 2)
      )
      assert result == [()]
      _, result, _ = where_op.infer_shape_partial(
          cond=(0,), x=(2, 2), y=(2, 2)
      )
      assert result == [()]
      with mx.np_shape(True):
          _, result, _ = where_op.infer_shape_partial(
              cond=(-1, 2), x=(-1, 2), y=(-1, 2)
          )
          assert result == [None]
          _, result, _ = where_op.infer_shape_partial(
              cond=(-1,), x=(2, 2), y=(2, 2)
          )
          assert result == [None]
  
file_path: mxnet/tests/python/unittest/test_infer_shape.py
imports: import mxnet
  from common import models
  import pytest
label_desc: Use image node_id 16 for calling a global function with example usage: test_where_partial_shape() without return types