Create a constant inside a graph

Hey,
I’m trying to edit a graph representation (torch._C.Graph) via python API.
I want to add a node at the begging of the graph that returns a CUDA device.

this is what I did:

cuda_node = graph.create('prim::Constant[value="cuda:0"]')
cuda_value = next(cuda_node.outputs())
cuda_value.setType(torch._C.DeviceObjType.get())
graph.prependNode(cuda_node)

This code has been successfully ran, but I get this error when I run the graph function by torch._C._create_function_from_graph(...) :

RuntimeError: 0 INTERNAL ASSERT FAILED at "/pytorch/torch/csrc/jit/ir/alias_analysis.cpp":465, please report a bug to PyTorch. We don't have an op for prim::Constant[value="cuda:0"] but it isn't a special case.  Argument types: 

How can I create this kind of node without experience this issue ?
Thanks.

It seems you are using internal methods from the _C namespace, so I’m unsure if your use case is supported out of the box (without breaking stuff).
What’s your use case you need to manipulate the graph and cannot recreate it?

@ptrblck I’m not sure that recreate a graph replica will help in that case, since that I still don’t know how to create a prim::Constant with the attribute value="cuda:0".

my mistake was that there is no prim::Constant[value="cuda:0"] node kind, I need to create the node by:

cuda_node = graph.create('prim::Constant')

and then I need to set for cuda_node an attribute value="cuda:0" but I’m not sure how.

This code snippet would create a prim::Constant on the GPU:

@torch.jit.script
def fun(x):
    y = x + torch.tensor(1, device='cuda')
    return y

print(fun.graph)
> graph(%x.1 : Tensor):
  %7 : bool = prim::Constant[value=0]()
  %21 : Device = prim::Constant[value="cuda"]()
  %5 : None = prim::Constant()
  %2 : int = prim::Constant[value=1]() # <ipython-input-36-11fdcf9e4060>:9:25
  %8 : Tensor = aten::tensor(%2, %5, %21, %7) # <ipython-input-36-11fdcf9e4060>:9:12
  %y.1 : Tensor = aten::add(%x.1, %8, %2) # <ipython-input-36-11fdcf9e4060>:9:8
  return (%y.1)

@ptrblck thanks, seems like I can take this example and use createClone to create my desired node in my graph. I guess there is no exposed API to set attributes for nodes, right?

I don’t know, so lets wait for some experts here. :wink:

@ptrblck actually, it takes too much time that I can’t afford to copy an entire graph just for 1 node modification. if there will be an option to set attributes for nodes it would be great.