Replace usage of xla.abstractify with core.get_aval where feasible. - #284

Replace usage of xla.abstractify with core.get_aval where feasible.#284

copybara-service[bot] wants to merge 1 commit into

maingoogle-deepmind/distrax:mainfrom

test_730598559google-deepmind/distrax:test_730598559Copy head branch name to clipboard