Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions QEfficient/customop/ctx_scatter_gather.py
Original file line number Diff line number Diff line change
Expand Up @@ -56,6 +56,7 @@ def symbolic(g: torch.Graph, data: torch.Value, position_ids: torch.Value, updat
return g.onnxscript_op(CtxScatter, data, position_ids, updates).setTypeAs(data)


# TODO : This function will be deprecated in future.
@onnxscript.script(onnxscript.values.Opset("com.qualcomm.cloud", 1))
def CtxScatter3D(data: onnxscript.FLOAT, position_ids: onnxscript.INT32, updates: onnxscript.FLOAT) -> onnxscript.FLOAT:
# Find dims
Expand All @@ -75,6 +76,7 @@ def CtxScatter3D(data: onnxscript.FLOAT, position_ids: onnxscript.INT32, updates
return ops.ScatterND(data, indices, updates)


# TODO : This function will be deprecated in future.
class CtxScatterFunc3D(torch.autograd.Function):
@staticmethod
def forward(data: torch.Tensor, position_ids: torch.Tensor, updates: torch.Tensor):
Expand All @@ -92,13 +94,15 @@ def symbolic(g: torch.Graph, data: torch.Value, position_ids: torch.Value, updat
return g.onnxscript_op(CtxScatter3D, data, position_ids, updates).setTypeAs(data)


# TODO : This function will be deprecated in future.
@onnxscript.script(onnxscript.values.Opset("com.qualcomm.cloud", 1))
def CtxGather3D(data: onnxscript.FLOAT, ctx_indices: onnxscript.INT32) -> onnxscript.FLOAT:
ctx_indices = ops.Expand(ctx_indices, ops.Slice(ops.Shape(data), starts=[0], ends=[2], axes=[0]))
ctx_indices = ops.Unsqueeze(ctx_indices, [-1])
return ops.GatherND(data, ctx_indices, batch_dims=1)


# TODO : This function will be deprecated in future.
class CtxGatherFunc3D(torch.autograd.Function):
@staticmethod
def forward(data: torch.Tensor, ctx_indices: torch.Tensor):
Expand Down
4 changes: 4 additions & 0 deletions QEfficient/customop/ctx_scatter_gather_cb.py
Original file line number Diff line number Diff line change
Expand Up @@ -56,6 +56,7 @@ def symbolic(
return g.onnxscript_op(CtxScatterCB, data, batch_index, position_ids, updates).setTypeAs(data)


# TODO : This function will be deprecated in future.
@onnxscript.script(onnxscript.values.Opset("com.qualcomm.cloud", 1))
def CtxScatterCB3D(
data: onnxscript.FLOAT, batch_index: onnxscript.INT32, position_ids: onnxscript.INT32, updates: onnxscript.FLOAT
Expand All @@ -76,6 +77,7 @@ def CtxScatterCB3D(
return ops.ScatterND(data, indices, updates)


# TODO : This function will be deprecated in future.
class CtxScatterFuncCB3D(torch.autograd.Function):
@staticmethod
def forward(data: torch.Tensor, batch_index: torch.Tensor, position_ids: torch.Tensor, updates: torch.Tensor):
Expand Down Expand Up @@ -139,6 +141,7 @@ def symbolic(
return g.onnxscript_op(CtxGatherCB, data, batch_index, ctx_indices, comp_ctx_len).setTypeAs(data)


# TODO : This function will be deprecated in future.
@onnxscript.script(onnxscript.values.Opset("com.qualcomm.cloud", 1))
def CtxGatherCB3D(
data: onnxscript.FLOAT, batch_index: onnxscript.INT32, ctx_indices: onnxscript.INT32
Expand All @@ -158,6 +161,7 @@ def CtxGatherCB3D(
return ops.GatherND(data, indices)


# TODO : This function will be deprecated in future.
class CtxGatherFuncCB3D(torch.autograd.Function):
@staticmethod
def forward(data: torch.Tensor, batch_index: torch.Tensor, ctx_indices: torch.Tensor):
Expand Down
Loading
Loading