Fixed a couple bugs with `pathwaysutils.jax`#119

copybara-service[bot] merged 1 commit intomainAI-Hypercomputer/pathways-utils:mainfrom

test_813906532AI-Hypercomputer/pathways-utils:test_813906532Copy head branch name to clipboard