Skip to content

Implement Mamba using EasyDel for JAX#448

Draft
ctr-pmuruganTT wants to merge 4 commits intomainfrom
pmurugan/implement_easydel_for_mamba
Draft

Implement Mamba using EasyDel for JAX#448
ctr-pmuruganTT wants to merge 4 commits intomainfrom
pmurugan/implement_easydel_for_mamba

Conversation

@ctr-pmuruganTT
Copy link
Contributor

Add EasyDel implementation for JAX Mamba model

What's changed

  • updated requirements.txt with the requirements for this model
  • updated load_model() to use EasyDel model
  • updated load_inputs() to pass mesh as a parameter
  • added load_parameters_partition_spec() and get_input_activations

Checklist

  • New/Existing tests provide coverage for changes

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants