fix jax backend set item with scalar values

[mattbarrett98]

this fixes the jax backend for set_item when scalar values are used i.e.

import ivy
ivy.set_jax_backend()
x = ivy.array([[1., 2.], [3., 4.]])
query = ivy.array([[True, False], [True, True]])
val = 5.
print(ivy.set_item(x, query, val))

[Ishticode]

hey @mattbarrett98 . I think there are cases where expected_shape is undefined. lmc if can be fixed unless you have something
you will see in this case

import ivy
ivy.set_jax_backend()
x = ivy.array([[1., 2.], [3., 4.]])
query = ivy.array([[1, 0], [1, 1]])
val = ivy.array(5.)
print(ivy.set_item(x, query, val))

also the test case is linked below

[Ishticode]

leads to expected_shape being undefined in one test case (https://github.com/unifyai/ivy/actions/runs/7463291702/job/20307666791?pr=27886).
Indenting the new conditional seems to work but not sure entirely.

  if ivy.is_array(query) and ivy.is_bool_dtype(query):
      query, expected_shape = _mask_to_index(query, x)
      if ivy.is_array(val):
          val = _broadcast_to(val, expected_shape)._data