Fix issue 389 vmap scalar assertions

[Sikandar1310291]

Summary

This PR fixes issue #389 by adding vmap compatibility to assert_scalar_positive, assert_scalar_non_negative, and assert_scalar_negative.

Changes

• Converted three scalar assertions from static to value assertions
• Added jittable predicate functions using checkify.check
• These assertions can now be used with jax.vmap, jax.jit, and jax.pmap when wrapped with @chex.chexify
• Added comprehensive test suite with 9 tests

Example (from issue #389)

import jax.numpy as jnp
import jax
from chex import assert_scalar_positive, chexify

x_vector = jnp.array([1., 1.])

@chexify
@jax.jit
def test_vmap():
    jax.vmap(assert_scalar_positive)(x_vector)
    return x_vector

test_vmap()  # Now works!