We pass in "static" objects like model, tx, config as JAX compares them by object hash, and instantiating them inside causes pjit top-level annotations to fail to match as pytree prefixes if we re-instantiate. Args: model, tx, config, is_training, key
(model, tx, config, is_training, key)
| 753 | |
| 754 | |
| 755 | def init_initial_state(model, tx, config, is_training, key): |
| 756 | """ |
| 757 | We pass in "static" objects like model, tx, config as JAX compares them by |
| 758 | object hash, and instantiating them inside causes pjit top-level annotations |
| 759 | to fail to match as pytree prefixes if we re-instantiate. |
| 760 | |
| 761 | Args: model, tx, config, is_training, key |
| 762 | """ |
| 763 | input_shape = (config.micro_batch_size_to_train_on, config.max_target_length) |
| 764 | image_shape = multimodal_utils.get_dummy_image_shape_for_init( |
| 765 | config.model_name, batch_size=config.micro_batch_size_to_train_on |
| 766 | ) |
| 767 | model_vars = model.init( |
| 768 | {"params": key, "dropout": key, "aqt": key}, |
| 769 | np.ones(input_shape, dtype=jnp.int32), |
| 770 | np.ones(input_shape, dtype=jnp.int32), |
| 771 | encoder_images=np.ones(image_shape, dtype=jnp.int32) if config.use_multimodal else None, |
| 772 | # nnx_method="no_op", |
| 773 | ) |
| 774 | if is_training: |
| 775 | return init_training_state(model.apply, model_vars, tx) |
| 776 | return init_decode_state(model.apply, model_vars) |
| 777 | |
| 778 | |
| 779 | def get_abstract_param(model, config): |
nothing calls this directly
no test coverage detected