d_hidden=16,
n_output_classes=n_classes)
inputs = lambda _: test_inputs(n_classes)
# Train and evaluate
state = trax.train(output_dir,
model=model_fn,
inputs=inputs,
train_steps=train_steps,
eval_steps=eval_steps,
optimizer=trax_opt.SM3)
# Assert total train steps
self.assertEqual(train_steps, state.step)
# Assert 2 evaluations ran
train_acc = state.history.get("train", "metrics/accuracy")
eval_acc = state.history.get("eval", "metrics/accuracy")
self.assertEqual(len(train_acc), len(eval_acc))
self.assertEqual(2, len(eval_acc))
# Predict with final params
inputs = inputs(1).train_stream()
model = layers.Serial(model_fn())
model(next(inputs)[0], state.params[0])
if __name__ == "__main__":
config.config_with_absl()
test.main()
def run(**executable_dict):
# JAX uses a different flags library -- parse from absl so that the task
# knows its task_id.
jax_config.config_with_absl()
app.run(functools.partial(main, executable_dict))
