cross_entropy calls soft_max https://github.com/pytorch/pytorch/blob/master/torch/nn/functional.py#L1645 so it would require couple line python change.

Yep, but you are already changing log_softmax, right?

@ssnl, Yes, I can do that.
Test failure is legit,

 ======================================================================
22:32:33 FAIL: test_passing_one_positional_but_not_the_second (__main__.TestCustomOperators)
22:32:33 ----------------------------------------------------------------------
22:32:33 RuntimeError: Found 2 overloads for operator aten::log_softmax! Overloads are not supported from Python.
22:32:33 
22:32:33 During handling of the above exception, another exception occurred:
22:32:33 
22:32:33 Traceback (most recent call last):
22:32:33   File "test_jit.py", line 7796, in test_passing_one_positional_but_not_the_second
22:32:33     torch.ops.aten.log_softmax(torch.ones(5))
22:32:33 AssertionError: "aten::log_softmax\(\) is missing value for argument 'dim'." does not match "Found 2 overloads for operator aten::log_softmax! Overloads are not supported from Python."

but I'm not sure what's the preferred fix should be. FWIW, some operators already have overloads that are not supported from python, e.g.

In [3]: torch.ops.aten.sum(torch.ones(5))
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-3-acd80e61fcec> in <module>()
----> 1 torch.ops.aten.sum(torch.ones(5))

/workspace/ngimel/pytorch_upstream/torch/_ops.py in __getattr__(self, op_name)
     56         # for overloads and raise an exception if there are more than one.
     57         qualified_op_name = '{}::{}'.format(self.name, op_name)
---> 58         op = torch._C._jit_get_operation(qualified_op_name)
     59         # let the script frontend know that op is identical to the builtin op
     60         # with qualified_op_name

RuntimeError: Found 5 overloads for operator aten::sum! Overloads are not supported from Python.

so log_softmax erroring out with similar message is not necessarily a big problem (?) .

cc @apaszke @jamesr66a on JIT test

@ngimel Looking at this more closely I would advise updating the error message here.

@ngimel Looking at this more closely I would advise updating the error message here.

In the jit tests on in the softmax_cpu assert? softmax_cpu should never be called with upconvert = True, and upconvert is not user exposed, if it happend, there's something wrong with the core that user can't fix, hence AT_ASSERTM and not AT_CHECK.

Anything I can do to help move this forward? @apaszke @ezyang

I'm not super happy with the upconvert flag. It doesn't really specify the destination type. Should it be float? Should it be double? The context is probably dependent on the device, and this seems to overfit the CUDA context. Can't we apply a simple modification to our kernels, or simply have a _log_softmax_half_to_float implemented only for CUDA, and dispatch to _log_softmax(...).to(dtype) otherwise?

upconvert is true only for cuda half inputs with fp32 dtype argument. I could dispatch to _log_softmax_half_to_float in this case, but it would require it's own entries for forward and backward in native_functions and in derivatives, and overall I don't think it would be any prettier.
Adding modification to kernels to support more input type /dtype combinations with a fast path can be done (in fact, for cuda kernels output type can be anything, it's a separate template parameter), but then dispatch will have to be really tricky (right now dispatch defines scalar_t from which I can derive acc_type, but any other combinations of input/output types would require changes to types defined in dispatch, and instantiating a cross-product of kernels with different input/output types, which no one wants.)

upconvert is true only for cuda half inputs with fp32 dtype argument

That's precisely the problem. It's a very specific flag, with a very specific meaning, which it not at all implied by its name/function name/function signature.

I don't understand why the dispatch would be a problem. Can't you just declare the derivatives for the top-level native function log_softmax, and have it take full responsibility for providing the derivative no matter which implementation it chooses?

If derivatives are declared for log_softmax then backward will have to take care of type conversion which is now delegated to autograd, which is not the end of the world, but will make backward more error-prone, especially if at some point some implicit conversions for other types are added.

Test failures seem unrelated.