Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Adjoint Refactor Phase 2 #1549

Open
16 tasks
tylerflex opened this issue Mar 15, 2024 · 0 comments
Open
16 tasks

Adjoint Refactor Phase 2 #1549

tylerflex opened this issue Mar 15, 2024 · 0 comments
Assignees
Labels
2.7 will go into version 2.7.*

Comments

@tylerflex
Copy link
Collaborator

Why?

  • Fix issues where JaxX methods inherited from td.X components ignore input_structures, output_monitors eg. JaxSimulation.epsilon ignores input_structures #1459
  • Simplify internal handling of Jax-traced structures and monitors
    • simpler accounting for how many of each type of monitor, structure.
    • simplify conversion from JaxSimulation to Simulation

Goals

  • JaxSimulation methods inherited from Simulation work without patching (eg .plot())
  • make web API support JaxSimulation upload directly.
  • eliminate need for separate JaxInfo file upload.

Strategy

  • Validate input_structures to make a td.Structure copy and append it to JaxSimulation.structures.
  • Validate output_monitors to make a td.Monitor copy and append it to JaxSimulation.monitors.
  • Similarly, add MonitorData from output_monitors to the JaxSimulationData.data tuple so it's accessible there.
  • Remove patching for JaxSimulationMethods
  • Simplify JaxSimulation.to_simulation() (just remove extraneous fields eg. output_monitors and input_structures and type-cast).
  • Simplify conversion from JaxSimulationData to adjoint JaxSimulation.
  • Make JaxSimulation upload able through web API, simplify the tda.web internals.
  • Remove JaxInfo from code, if possible.

If possible.

  • Remove need for separate JaxStructureStaticX, allow mixing of jax and regular types.
@tylerflex tylerflex added adjoint 2.7 will go into version 2.7.* labels Mar 15, 2024
@tylerflex tylerflex self-assigned this Mar 15, 2024
@tylerflex tylerflex removed the adjoint label May 24, 2024
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
2.7 will go into version 2.7.*
Projects
None yet
Development

No branches or pull requests

1 participant