[jax-inference-offloading] consolidate definitions for default tensor dtype#1816

yhtang wants to merge 6 commits intomainNVIDIA/JAX-Toolbox:mainfrom

yhtang/jio-default-dtypeNVIDIA/JAX-Toolbox:yhtang/jio-default-dtypeCopy head branch name to clipboard