Add device/import utilities, expose vectorize/devectorize methods as public callables - #83
Conversation
anotheranshu
left a comment
anotheranshu
left a comment
There was a problem hiding this comment.
None of the comments are blocking
| """ | ||
| if shutil.which("nvcc") is not None: | ||
| return True | ||
| return CUDA_HOME is not None and os.path.isfile(os.path.join(CUDA_HOME, "bin", "nvcc")) |
There was a problem hiding this comment.
Do we need/want this fallback, or should we assume that users have CUDA_HOME on their PATH?
There was a problem hiding this comment.
CUDA_HOME is a commonly used env var that is set when cuda /cuda nvcc gets installed, and I think it's common for it to not be in PATH which shutil looks for (but checking for both allows us to cover all bases).
| sample_weight = None | ||
| if block_sensitivity is not None: | ||
| sens_vectorized = self._vectorize(block_sensitivity) | ||
| sens_vectorized = self._vectorize_block(block_sensitivity) |
There was a problem hiding this comment.
nit (and not really this PR): can we rename sens_ to sensitivity_?
| """ | ||
| weight_shape, weight_dtype = weight.shape, weight.dtype | ||
| blocks, axis = self._scale_reshape_and_block(weight) | ||
| block_shape = blocks[0].shape # all blocks share one shape |
There was a problem hiding this comment.
Should we validate this and error out here if there are misshapen blocks?
There was a problem hiding this comment.
Oh hmm nvm we do verify this deep down the call stack of _scale_reshape_and_block. I am a little concerned that we're relying on the check so far away, but given that we aren't forking to multiple implementations under the hood I think it's fine
2 participants
Uh oh!
There was an error while loading. Please reload this page.