MCPcopy Create free account
hub / github.com/apache/tvm / init

Method init

python/tvm/relax/training/optimizer.py:275–302  ·  view source on GitHub ↗

Set the parameters, determine the dtype, and construct the initial state for the optimizer. The state of SGD is `(num_steps,)`. Parameters ---------- params : Union[Var, List[Var]] The parameter or the list of parameters to optimize.

(self, params: Var | list[Var])

Source from the content-addressed store, hash-verified

273 self.weight_decay = float(weight_decay)
274
275 def init(self, params: Var | list[Var]) -> "SGD":
276 """Set the parameters, determine the dtype, and construct the initial state for the
277 optimizer.
278
279 The state of SGD is `(num_steps,)`.
280
281 Parameters
282 ----------
283 params : Union[Var, List[Var]]
284 The parameter or the list of parameters to optimize.
285
286 Parameters should all be Vars of floating point Tensors, including float32, float64,
287 float16, etc. Currently, all parameters should have the same dtype, and that dtype
288 will be used as the dtype of the optimizer states.
289
290 Returns
291 -------
292 self : SGD
293 The SGD optimizer itself.
294 """
295 if not isinstance(params, list):
296 params = [params]
297 self._set_params_and_dtype(params)
298 self.state = (
299 # num_steps = 0
300 tvm.runtime.tensor(np.zeros((), "int64")),
301 )
302 return self
303
304 def get_function(self) -> Function:
305 """Use blockbuilder to construct an optimizer function that executes updates of the

Callers 5

batch_tree_attnFunction · 0.45
tree_attn_paged_kvFunction · 0.45
compute_s_gemmFunction · 0.45
compute_o_gemmFunction · 0.45
transform_moduleMethod · 0.45

Calls 2

_set_params_and_dtypeMethod · 0.80
zerosMethod · 0.45

Tested by

no test coverage detected