Add more type support for burn-jit #2454
Conversation
| SequenceArg::new(), | ||
| SequenceArg::new(), | ||
| SequenceArg::new(), | ||
| SequenceArg::new(), | ||
| SequenceArg::new(), | ||
| SequenceArg::new(), | ||
| SequenceArg::new(), | ||
| SequenceArg::new(), | ||
| SequenceArg::new(), | ||
| SequenceArg::new(), | ||
| SequenceArg::new(), | ||
| SequenceArg::new(), |
There was a problem hiding this comment.
Not sure if there is a better way to do this 😅
There was a problem hiding this comment.
Tbh I don't understand fusion at all, so I just copied what was there for all the new types 😅
There was a problem hiding this comment.
It's simply that each precision has its state. It would be cleaner to have a map Map<ElemenType, State>, instead of listing all posibilities. I'll think how I can refactor it, but not in this PR.
| use super::*; | ||
| use burn_jit::tests::{burn_autodiff, burn_fusion, burn_ndarray, burn_tensor}; | ||
|
|
||
| pub type TestBackend = burn_fusion::Fusion<JitBackend<TestRuntime, f32, i32>>; | ||
| pub type TestBackend2<F, I> = burn_fusion::Fusion<JitBackend<TestRuntime, F, I>>; |
There was a problem hiding this comment.
TestBackend2 is the one that can be injected with the precision right? Not sure if we could come up with a better name.
There was a problem hiding this comment.
I couldn't come up with a better name and it's only used in the macro to instantiate and override TestBackend with the correct types, so I left it for now. It's not used in the actual tests, just the macro.
Pull Request Template
Checklist
run-checks allscript has been executed.Related Issues/PRs
Implements the new types and parameterized testing from tracel-ai/cubecl#207
Changes
Adds support for
flex32,i8,i16,i64,u8,u16,u64support to burn, parameterize tests to ensure all types work correctlyTesting
New parameterized tests