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

Add parameter freezing STOLayer and general performance fix #67

Merged
merged 24 commits into from
Oct 8, 2023

Conversation

CheukHinHoJerry
Copy link
Collaborator

@CheukHinHoJerry CheukHinHoJerry commented Sep 22, 2023

Updates from this branch (To be updated)

  • add freeze parameter STOLayer
  • type stability fix on some basis
  • commented out FlexArray interface
  • skip bounds check in SparseSymProd inside closure
  • prevent calling generic_matmul! in rrule of LinearLayer

cc @DexuanZhou

@cortner : will let you know once it is ready.

@CheukHinHoJerry CheukHinHoJerry changed the title Benchmarks and performance fix Add parameter freezing STOLayer and general performance fix Sep 29, 2023
src/lux.jl Outdated
@@ -63,6 +63,7 @@ evaluate!(out, basis::AbstractPoly4MLBasis, X, ps) = evaluate!(out, basis, X)
function evaluate(l::PolyLuxLayer, X, ps, st)
out = acquire!(st.pool, _outsym(X), _out_size(l.basis, X), _valtype(l.basis, X))
evaluate!(out, l.basis, X, ps)
release!(X)
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

flag in LuxLayer on whether or not inputs can be released.

@cortner
Copy link
Member

cortner commented Oct 8, 2023

@CheukHinHoJerry -- the only change since we discussed it is the flag for whether or not to release? I.e. I should just go ahead and merge this?

@CheukHinHoJerry
Copy link
Collaborator Author

@CheukHinHoJerry -- the only change since we discussed it is the flag for whether or not to release? I.e. I should just go ahead and merge this?

Yes - I think this is the only thing. And we talked about moving AtomicorbitalRadius and wf related stuffs out of this package but I think it should not be done in this PR anyways.

@cortner cortner merged commit f55e4b9 into main Oct 8, 2023
3 of 4 checks passed
@cortner cortner deleted the benchmarks branch October 8, 2023 20:57
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.

3 participants