Method batch

axlearn/common/utils_test.py:651–652 · view source on GitHub ↗

(fn)

Source from the content-addressed store, hash-verified

649	key, child_key = jax.random.split(key)
650	base_results.append(fn(child_key))
651	base_results = jnp.stack(base_results)
652
653	def batch(fn):
654	return lambda split_keys: jax.vmap(fn)(split_keys.keys)
655

dataset_fnFunction · 0.80

_asr_text_datasetFunction · 0.80

get_dataset_cardinalityFunction · 0.80

test_fake_multiple_choice_postprocessMethod · 0.80

test_dispatch_tpuMethod · 0.80

per_feed_batchFunction · 0.80

processFunction · 0.80

_pad_logical_to_physicalFunction · 0.80

fnFunction · 0.80

_mask_fnMethod · 0.80

process_dataset_fnFunction · 0.80

windowed_packingFunction · 0.80

no outgoing calls

no test coverage detected