iree: INVALID_ARGUMENT; bindings[1] is indirect but the command buffer does not support binding tables

What happened?

I am trying to run a StableHLO model exported from JAX. It compiles without errors, but then at runtime I get:

  File "/home/petkantchin/ws/iree/build/ninja/Debug/runtime/bindings/python/iree/runtime/function.py", line 137, in __call__
    self._invoke(arg_list, ret_list)
  File "/home/petkantchin/ws/iree/build/ninja/Debug/runtime/bindings/python/iree/runtime/function.py", line 162, in _invoke
    self._vm_context.invoke(self._vm_function, arg_list, ret_list)
ValueError: Error invoking function: repo/runtime/src/iree/hal/command_buffer_validation.c:514: INVALID_ARGUMENT; bindings[1] is indirect but the command buffer does not support binding tables; while invoking native function hal.command_buffer.push_descriptor_set; while calling import; 
[ 1]   native hal.command_buffer.push_descriptor_set:0 -
[ 0] bytecode jit_train_step.main:92338 [
    stable_diffusion_train_step.mlir:695:12
      at stable_diffusion_train_step.mlir:2:3,
    stable_diffusion_train_step.mlir:16518:10
      at stable_diffusion_train_step.mlir:16468:10
      at stable_diffusion_train_step.mlir:693:12
      at stable_diffusion_train_step.mlir:2:3,
    stable_diffusion_train_step.mlir:16500:11
      at stable_diffusion_train_step.mlir:693:12
      at stable_diffusion_train_step.mlir:2:3,
    stable_diffusion_train_step.mlir:16509:27
      at stable_diffusion_train_step.mlir:693:12
      at stable_diffusion_train_step.mlir:2:3,
    stable_diffusion_train_step.mlir:16510:11
      at stable_diffusion_train_step.mlir:693:12
      at stable_diffusion_train_step.mlir:2:3,
    stable_diffusion_train_step.mlir:16574:11
      at stable_diffusion_train_step.mlir:694:12
      at stable_diffusion_train_step.mlir:2:3,
    stable_diffusion_train_step.mlir:16578:10
      at stable_diffusion_train_step.mlir:16528:10
      at stable_diffusion_train_step.mlir:694:12
      at stable_diffusion_train_step.mlir:2:3,
    stable_diffusion_train_step.mlir:16463:10
      at stable_diffusion_train_step.mlir:693:12
      at stable_diffusion_train_step.mlir:2:3,
    stable_diffusion_train_step.mlir:16560:11
      at stable_diffusion_train_step.mlir:694:12
      at stable_diffusion_train_step.mlir:2:3,
    stable_diffusion_train_step.mlir:16570:11
      at stable_diffusion_train_step.mlir:694:12
      at stable_diffusion_train_step.mlir:2:3,
    stable_diffusion_train_step.mlir:16527:10
      at stable_diffusion_train_step.mlir:694:12
      at stable_diffusion_train_step.mlir:2:3,
    stable_diffusion_train_step.mlir:2:3,
    stable_diffusion_train_step.mlir:16467:10
      at stable_diffusion_train_step.mlir:693:12
      at stable_diffusion_train_step.mlir:2:3
  ]

Steps to reproduce your issue

Download iree_no_binding_table_bug.zip, unzip it.

cd iree_no_binding_table_bug
python iree_no_binding_table_bug.py

What component(s) does this issue relate to?

Python, Runtime

Version information

I had to fix another bug related to 0-rank tensors since some of the inputs are of this form. The PR for this is pending at https://github.com/openxla/iree/pull/14298.

Additional context

No response